package cn.cloud.all.security.crypto.argon2;

import java.util.Base64;

import org.bouncycastle.crypto.params.Argon2Parameters;
import org.bouncycastle.util.Arrays;

class Argon2EncodingUtils {
    private static final Base64.Encoder b64encoder = Base64.getEncoder().withoutPadding();
    private static final Base64.Decoder b64decoder = Base64.getDecoder();

    public static String encode(byte[] hash, Argon2Parameters parameters) throws IllegalArgumentException {
        StringBuilder stringBuilder = new StringBuilder();

        switch (parameters.getType()) {
            case Argon2Parameters.ARGON2_d:
                stringBuilder.append("$argon2d");
                break;
            case Argon2Parameters.ARGON2_i:
                stringBuilder.append("$argon2i");
                break;
            case Argon2Parameters.ARGON2_id:
                stringBuilder.append("$argon2id");
                break;
            default:
                throw new IllegalArgumentException("Invalid algorithm type: " + parameters.getType());
        }
        stringBuilder.append("$v=").append(parameters.getVersion())
                .append("$m=").append(parameters.getMemory())
                .append(",t=").append(parameters.getIterations())
                .append(",p=").append(parameters.getLanes());

        if (parameters.getSalt() != null) {
            stringBuilder.append("$").append(b64encoder.encodeToString(parameters.getSalt()));
        }

        stringBuilder.append("$").append(b64encoder.encodeToString(hash));

        return stringBuilder.toString();
    }

    public static Argon2Hash decode(String encodedHash) throws IllegalArgumentException {
        Argon2Parameters.Builder paramsBuilder;

        String[] parts = encodedHash.split("\\$");

        if (parts.length < 4) {
            throw new IllegalArgumentException("Invalid encoded Argon2-hash");
        }

        int currentPart = 1;

        switch (parts[currentPart++]) {
            case "argon2d":
                paramsBuilder = new Argon2Parameters.Builder(Argon2Parameters.ARGON2_d);
                break;
            case "argon2i":
                paramsBuilder = new Argon2Parameters.Builder(Argon2Parameters.ARGON2_i);
                break;
            case "argon2id":
                paramsBuilder = new Argon2Parameters.Builder(Argon2Parameters.ARGON2_id);
                break;
            default:
                throw new IllegalArgumentException("Invalid algorithm type: " + parts[0]);
        }

        if (parts[currentPart].startsWith("v=")) {
            paramsBuilder.withVersion(Integer.parseInt(parts[currentPart].substring(2)));
            currentPart++;
        }

        String[] performanceParams = parts[currentPart++].split(",");

        if (performanceParams.length != 3) {
            throw new IllegalArgumentException("Amount of performance parameters invalid");
        }

        if (performanceParams[0].startsWith("m=")) {
            paramsBuilder.withMemoryAsKB(Integer.parseInt(performanceParams[0].substring(2)));
        } else {
            throw new IllegalArgumentException("Invalid memory parameter");
        }

        if (performanceParams[1].startsWith("t=")) {
            paramsBuilder.withIterations(Integer.parseInt(performanceParams[1].substring(2)));
        } else {
            throw new IllegalArgumentException("Invalid iterations parameter");
        }

        if (performanceParams[2].startsWith("p=")) {
            paramsBuilder.withParallelism(Integer.parseInt(performanceParams[2].substring(2)));
        } else {
            throw new IllegalArgumentException("Invalid parallelity parameter");
        }

        paramsBuilder.withSalt(b64decoder.decode(parts[currentPart++]));

        return new Argon2Hash(b64decoder.decode(parts[currentPart]), paramsBuilder.build());
    }

    public static class Argon2Hash {

        private byte[] hash;
        private Argon2Parameters parameters;

        Argon2Hash(byte[] hash, Argon2Parameters parameters) {
            this.hash = Arrays.clone(hash);
            this.parameters = parameters;
        }

        public byte[] getHash() {
            return Arrays.clone(hash);
        }

        public void setHash(byte[] hash) {
            this.hash = Arrays.clone(hash);
        }

        public Argon2Parameters getParameters() {
            return parameters;
        }

        public void setParameters(Argon2Parameters parameters) {
            this.parameters = parameters;
        }
    }
}
